# 本程序 计划弄清dolfinx程序变量的值的操作
# 关注重点是 刚度矩阵和载荷向量 
# Demo
# u(x, y) = sin(x)cos(y) f(x, y) = 2sin(x)cos(y)
# -Laplace(u) = f in \Omega
# u = g on \partial\Omega
import importlib.util

if importlib.util.find_spec("petsc4py") is not None:
    import dolfinx

    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
    from petsc4py.PETSc import ScalarType  # type: ignore
else:
    print("This demo requires petsc4py.")
    exit(0)

from mpi4py import MPI

import numpy as np
from petsc4py import PETSc
import ufl
from dolfinx import fem, io, mesh, plot
from dolfinx.fem.petsc import LinearProblem, assemble_matrix, assemble_vector, create_matrix
from ufl import ds, dx, grad, inner
from dolfinx.fem.petsc import set_bc  # 设置边界条件新api

def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    # 需要导入
    # from petsc4py import PETSc
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data


# +
msh = mesh.create_rectangle(
    comm=MPI.COMM_WORLD,
    points=((0.0, 0.0), (1.0, 1.0)),
    n=(2, 2),
    cell_type=mesh.CellType.triangle,
)
V = fem.functionspace(msh, ("Lagrange", 1))


facets = mesh.locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 1.0) | np.isclose(x[1], 0.0) | np.isclose(x[1], 1.0)
)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
x = ufl.SpatialCoordinate(msh)
def u_true(x):
    return np.sin(x[0]) * np.cos(x[1])


f = 2 * ufl.sin(x[0]) * ufl.cos(x[1])
a = inner(grad(u), grad(v)) * dx  # <class 'ufl.form.Form'>
a_form = fem.form(a) # < class 'dolfinx.fem.forms.Form' >
L = inner(f, v) * dx
L_form = fem.form(L)

u_err = dolfinx.fem.Function(V)
u_exact = dolfinx.fem.Function(V)
u_exact.interpolate(u_true)

dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)
bc = fem.dirichletbc(u_exact, dofs=dofs)


A = assemble_matrix(a_form, [bc])   # 刚度矩阵
A.assemble()
b = assemble_vector(L_form)  # 载荷向量 需要手动处理自由度

fem.apply_lifting(b, [a_form], [[bc]])  # 调整载荷项b，考虑边界条件影响
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
set_bc(b, [bc])

A_mat = getStiffnessMatrix(A)  # ndarray
# 2*2剖分 A_max.shape=(9,9) b_vec.shape=(9,)
uh_np = np.linalg.solve(A_mat, b.array)  # np格式求解


problem = LinearProblem(a, L, bcs=[bc], petsc_options={
                        "ksp_type": "preonly", "pc_type": "lu"})
uh = problem.solve()

# 使用solver进行求解
solver = PETSc.KSP().create(msh.comm)
solver.setOperators(A)  # constant: A
solver.setType(PETSc.KSP.Type.PREONLY)
solver.getPC().setType(PETSc.PC.Type.LU)

u_solver = fem.Function(V)

solver.solve(b, u_solver.x.petsc_vec)

print(f"analytical solution:\n{u_exact.x.array}")
print(f"dolfinx solve:\n{uh.x.array}")
print(f"numpy solve:\n{uh_np}")  # 本质和下面是一样的 numpy转成了稠密矩阵运算
print(f"solver:\n{u_solver.x.array}")  # 这个是官方接口
